# 小题2：SigLIP zero-shot classification
import os
import glob
from PIL import Image
from transformers import pipeline
import re
# SigLIP模型准备
ckpt="D:/model/siglip2"
classifier=pipeline(model=ckpt,task="zero-shot-image-classification",device="cuda")
# 从README文件提取食物类别标签
with open("5/5.2/README (1).md","r",encoding="utf-8") as f:
    readme_content=f.read()
label_pattern=r"'(\d+)':\s*(\w+(?:_\w+)*)"
matches=re.findall(label_pattern,readme_content)
food_labels=[match[1].replace("_"," ") for match in sorted(matches,key=lambda x: int(x[0]))]
# 读取每个类别文件夹的前10张图片
photo_dir="D:/model/siglip2/photo/images"
selected_data=[]
for class_folder in os.listdir(photo_dir):
    class_path=f"{photo_dir}/{class_folder}"
    if os.path.isdir(class_path):
        image_files=glob.glob(f"{class_path}/*.jpg")[:10]
        for image_file in image_files:
            selected_data.append({"image_path":image_file,"class_name":class_folder})
# 执行zero-shot分类
correct_predictions=0
total_images=len(selected_data)
for item in selected_data:
    image=Image.open(item["image_path"])
    true_label_text=item["class_name"].replace("_"," ")
    predictions=classifier(image,food_labels)
    top5_labels=[pred["label"] for pred in predictions[:5]]
    if true_label_text in top5_labels:
        correct_predictions+=1
# 计算准确率
top5_accuracy=correct_predictions/total_images
print(f"Top-5 Accuracy: {top5_accuracy:.4f}")
print(f"Correct predictions: {correct_predictions}/{total_images}")